{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  from ._conv import register_converters as _register_converters\n"
     ]
    }
   ],
   "source": [
    "#reader.py\n",
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import collections\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "Py3 = sys.version_info[0] == 3\n",
    "\n",
    "def _read_words(filename):\n",
    "  with tf.gfile.GFile(filename, \"r\") as f:\n",
    "    if Py3:\n",
    "      return f.read().replace(\"\\n\", \"<eos>\").split()\n",
    "    else:\n",
    "      return f.read().decode(\"utf-8\").replace(\"\\n\", \"<eos>\").split()\n",
    "\n",
    "\n",
    "def _build_vocab(filename):\n",
    "  data = _read_words(filename)\n",
    "\n",
    "  counter = collections.Counter(data)\n",
    "  count_pairs = sorted(counter.items(), key=lambda x: (-x[1], x[0]))\n",
    "\n",
    "  words, _ = list(zip(*count_pairs))\n",
    "  word_to_id = dict(zip(words, range(len(words))))\n",
    "\n",
    "  return word_to_id\n",
    "\n",
    "\n",
    "def _file_to_word_ids(filename, word_to_id):\n",
    "  data = _read_words(filename)\n",
    "  return [word_to_id[word] for word in data if word in word_to_id]\n",
    "\n",
    "\n",
    "def ptb_raw_data(data_path=None):\n",
    "  \"\"\"Load PTB raw data from data directory \"data_path\".\n",
    "  Reads PTB text files, converts strings to integer ids,\n",
    "  and performs mini-batching of the inputs.\n",
    "  The PTB dataset comes from Tomas Mikolov's webpage:\n",
    "  http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz\n",
    "  Args:\n",
    "    data_path: string path to the directory where simple-examples.tgz has\n",
    "      been extracted.\n",
    "  Returns:\n",
    "    tuple (train_data, valid_data, test_data, vocabulary)\n",
    "    where each of the data objects can be passed to PTBIterator.\n",
    "  \"\"\"\n",
    "\n",
    "  train_path = os.path.join(data_path, \"ptb.train.txt\")\n",
    "  valid_path = os.path.join(data_path, \"ptb.valid.txt\")\n",
    "  test_path = os.path.join(data_path, \"ptb.test.txt\")\n",
    "\n",
    "  word_to_id = _build_vocab(train_path)\n",
    "  train_data = _file_to_word_ids(train_path, word_to_id)\n",
    "  valid_data = _file_to_word_ids(valid_path, word_to_id)\n",
    "  test_data = _file_to_word_ids(test_path, word_to_id)\n",
    "  vocabulary = len(word_to_id)\n",
    "  return train_data, valid_data, test_data, vocabulary\n",
    "\n",
    "\n",
    "def ptb_producer(raw_data, batch_size, num_steps, name=None):\n",
    "  \"\"\"Iterate on the raw PTB data.\n",
    "  This chunks up raw_data into batches of examples and returns Tensors that\n",
    "  are drawn from these batches.\n",
    "  Args:\n",
    "    raw_data: one of the raw data outputs from ptb_raw_data.\n",
    "    batch_size: int, the batch size.\n",
    "    num_steps: int, the number of unrolls.\n",
    "    name: the name of this operation (optional).\n",
    "  Returns:\n",
    "    A pair of Tensors, each shaped [batch_size, num_steps]. The second element\n",
    "    of the tuple is the same data time-shifted to the right by one.\n",
    "  Raises:\n",
    "    tf.errors.InvalidArgumentError: if batch_size or num_steps are too high.\n",
    "  \"\"\"\n",
    "  with tf.name_scope(name, \"PTBProducer\", [raw_data, batch_size, num_steps]):\n",
    "    raw_data = tf.convert_to_tensor(raw_data, name=\"raw_data\", dtype=tf.int32)\n",
    "\n",
    "    data_len = tf.size(raw_data)\n",
    "    batch_len = data_len // batch_size\n",
    "    data = tf.reshape(raw_data[0 : batch_size * batch_len],\n",
    "                      [batch_size, batch_len])\n",
    "\n",
    "    epoch_size = (batch_len - 1) // num_steps\n",
    "    assertion = tf.assert_positive(\n",
    "        epoch_size,\n",
    "        message=\"epoch_size == 0, decrease batch_size or num_steps\")\n",
    "    with tf.control_dependencies([assertion]):\n",
    "      epoch_size = tf.identity(epoch_size, name=\"epoch_size\")\n",
    "\n",
    "    i = tf.train.range_input_producer(epoch_size, shuffle=False).dequeue()\n",
    "    x = tf.strided_slice(data, [0, i * num_steps],\n",
    "                         [batch_size, (i + 1) * num_steps])\n",
    "    x.set_shape([batch_size, num_steps])\n",
    "    y = tf.strided_slice(data, [0, i * num_steps + 1],\n",
    "                         [batch_size, (i + 1) * num_steps + 1])\n",
    "    y.set_shape([batch_size, num_steps])\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "# 运行上面的reader.py\n",
    "\n",
    "data_path = './data/PTB-dataset-Tomas-Mikolov/data'\n",
    "# 隐藏层单元数与LSTM层级数\n",
    "hidden_size = 200\n",
    "num_layers = 2\n",
    "#词典规模\n",
    "vocab_size = 10000\n",
    "\n",
    "learning_rate = 1.0\n",
    "train_batch_size = 16\n",
    "# 训练数据截断长度\n",
    "train_num_step = 32\n",
    "\n",
    "# 在测试时不需要使用截断，测试数据为一个超长序列\n",
    "eval_batch_size = 1\n",
    "eval_num_step = 1\n",
    "num_epoch = 3\n",
    "#结点不被Dropout的概率\n",
    "keep_prob = 0.5\n",
    "\n",
    "# 用于控制梯度爆炸的参数\n",
    "max_grad_norm = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 通过ptbmodel 的类描述模型\n",
    "class PTBModel(object):\n",
    "    def __init__(self, is_training, batch_size, num_steps):\n",
    "        # 记录使用的Batch大小和截断长度\n",
    "        self.batch_size = batch_size\n",
    "        self.num_steps = num_steps\n",
    "\n",
    "        # 定义输入层，维度为批量大小×截断长度\n",
    "        self.input_data = tf.placeholder(tf.int32, [batch_size, num_steps])\n",
    "        # 定义预期输出\n",
    "        self.targets = tf.placeholder(tf.int32, [batch_size, num_steps])\n",
    "\n",
    "        # 定义使用LSTM结构为循环体，带Dropout的深度RNN\n",
    "        lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(hidden_size)\n",
    "        if is_training:\n",
    "            lstm_cell = tf.nn.rnn_cell.DropoutWrapper(lstm_cell, output_keep_prob=keep_prob)\n",
    "        cell = tf.nn.rnn_cell.MultiRNNCell([lstm_cell] * num_layers)\n",
    "\n",
    "        # 初始化状态为0\n",
    "        self.initial_state = cell.zero_state(batch_size, tf.float32)\n",
    "\n",
    "        # 将单词ID转换为单词向量，embedding的维度为vocab_size*hidden_size\n",
    "        embedding = tf.get_variable('embedding', [vocab_size, hidden_size])\n",
    "        # 将一个批量内的单词ID转化为词向量，转化后的输入维度为批量大小×截断长度×隐藏单元数\n",
    "        inputs = tf.nn.embedding_lookup(embedding, self.input_data)\n",
    "\n",
    "        # 只在训练时使用Dropout\n",
    "        if is_training: inputs = tf.nn.dropout(inputs, keep_prob)\n",
    "\n",
    "        # 定义输出列表，这里先将不同时刻LSTM的输出收集起来，再通过全连接层得到最终输出\n",
    "        outputs = []\n",
    "        # state 储存不同批量中LSTM的状态，初始为0\n",
    "        state = self.initial_state\n",
    "        with tf.variable_scope('RNN'):\n",
    "            for time_step in range(num_steps):\n",
    "                if time_step > 0: tf.get_variable_scope().reuse_variables()\n",
    "                # 从输入数据获取当前时间步的输入与前一时间步的状态，并传入LSTM结构\n",
    "                cell_output, state = cell(inputs[:, time_step, :], state)\n",
    "                # 将当前输出加入输出队列\n",
    "                outputs.append(cell_output)\n",
    "\n",
    "        # 将输出队列展开成[batch,hidden*num_step]的形状，再reshape为[batch*num_step, hidden]\n",
    "        output = tf.reshape(tf.concat(outputs, 1), [-1, hidden_size])\n",
    "\n",
    "        # 将LSTM的输出传入全连接层以生成最后的预测结果。最后结果在每时刻上都是长度为vocab_size的张量\n",
    "        # 且经过softmax层后表示下一个位置不同词的概率\n",
    "        weight = tf.get_variable('weight', [hidden_size, vocab_size])\n",
    "        bias = tf.get_variable('bias', [vocab_size])\n",
    "        logits = tf.matmul(output, weight) + bias\n",
    "\n",
    "        # 定义交叉熵损失函数，一个序列的交叉熵之和\n",
    "        loss = tf.contrib.legacy_seq2seq.sequence_loss_by_example(\n",
    "            [logits],  # 预测的结果\n",
    "            [tf.reshape(self.targets, [-1])],  # 期望正确的结果，这里将[batch_size, num_steps]压缩为一维张量\n",
    "            [tf.ones([batch_size * num_steps], dtype=tf.float32)])  # 损失的权重，所有为1表明不同批量和时刻的重要程度一样\n",
    "\n",
    "        # 计算每个批量的平均损失\n",
    "        self.cost = tf.reduce_sum(loss) / batch_size\n",
    "        self.final_state = state\n",
    "\n",
    "        # 只在训练模型时定义反向传播操作\n",
    "        if not is_training: return\n",
    "        trainable_variable = tf.trainable_variables()\n",
    "\n",
    "        # 控制梯度爆炸问题\n",
    "        grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, trainable_variable), max_grad_norm)\n",
    "        # 如果需要使用Adam作为优化器，可以改为tf.train.AdamOptimizer(learning_rate)，学习率需要降低至0.001左右\n",
    "        optimizer = tf.train.GradientDescentOptimizer(learning_rate)\n",
    "        # 定义训练步骤\n",
    "        self.train_op = optimizer.apply_gradients(zip(grads, trainable_variable))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def run_epoch(session, model, data, train_op, output_log, epoch_size):\n",
    "    total_costs = 0.0\n",
    "    iters = 0\n",
    "    state = session.run(model.initial_state)\n",
    "\n",
    "    # # 使用当前数据训练或测试模型\n",
    "    for step in range(epoch_size):\n",
    "        x, y = session.run(data)\n",
    "        # 在当前批量上运行train_op并计算损失值，交叉熵计算的是下一个单词为给定单词的概率\n",
    "        cost, state, _ = session.run([model.cost, model.final_state, train_op],\n",
    "                                        {model.input_data: x, model.targets: y, model.initial_state: state})\n",
    "        # 将不同时刻和批量的概率就可得到困惑度的对数形式，将这个和做指数运算就可得到困惑度\n",
    "        total_costs += cost\n",
    "        iters += model.num_steps\n",
    "\n",
    "        # 只在训练时输出日志\n",
    "        if output_log and step % 100 == 0:\n",
    "            print(\"After %d steps, perplexity is %.3f\" % (step, np.exp(total_costs / iters)))\n",
    "    return np.exp(total_costs / iters)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In iteration: 1\n",
      "After 0 steps, perplexity is 9997.503\n",
      "After 100 steps, perplexity is 1430.310\n",
      "After 200 steps, perplexity is 1019.812\n",
      "After 300 steps, perplexity is 884.877\n",
      "After 400 steps, perplexity is 770.536\n",
      "After 500 steps, perplexity is 685.862\n",
      "After 600 steps, perplexity is 627.773\n",
      "After 700 steps, perplexity is 584.603\n",
      "After 800 steps, perplexity is 550.645\n",
      "After 900 steps, perplexity is 522.333\n",
      "After 1000 steps, perplexity is 497.754\n",
      "After 1100 steps, perplexity is 477.317\n",
      "After 1200 steps, perplexity is 459.825\n",
      "After 1300 steps, perplexity is 446.264\n",
      "After 1400 steps, perplexity is 433.608\n",
      "After 1500 steps, perplexity is 423.149\n",
      "After 1600 steps, perplexity is 409.957\n",
      "After 1700 steps, perplexity is 400.390\n",
      "After 1800 steps, perplexity is 393.291\n",
      "Epoch: 1 Validation Perplexity: 239.055\n",
      "In iteration: 2\n",
      "After 0 steps, perplexity is 422.337\n",
      "After 100 steps, perplexity is 285.401\n",
      "After 200 steps, perplexity is 266.674\n",
      "After 300 steps, perplexity is 273.085\n",
      "After 400 steps, perplexity is 267.213\n",
      "After 500 steps, perplexity is 257.644\n",
      "After 600 steps, perplexity is 252.940\n",
      "After 700 steps, perplexity is 249.936\n",
      "After 800 steps, perplexity is 248.188\n",
      "After 900 steps, perplexity is 246.797\n",
      "After 1000 steps, perplexity is 244.862\n",
      "After 1100 steps, perplexity is 243.292\n",
      "After 1200 steps, perplexity is 241.627\n",
      "After 1300 steps, perplexity is 241.469\n",
      "After 1400 steps, perplexity is 240.662\n",
      "After 1500 steps, perplexity is 240.259\n",
      "After 1600 steps, perplexity is 237.447\n",
      "After 1700 steps, perplexity is 236.469\n",
      "After 1800 steps, perplexity is 236.575\n",
      "Epoch: 2 Validation Perplexity: 196.142\n",
      "In iteration: 3\n",
      "After 0 steps, perplexity is 353.527\n",
      "After 100 steps, perplexity is 236.278\n",
      "After 200 steps, perplexity is 221.385\n",
      "After 300 steps, perplexity is 228.476\n",
      "After 400 steps, perplexity is 224.843\n",
      "After 500 steps, perplexity is 217.486\n",
      "After 600 steps, perplexity is 214.905\n",
      "After 700 steps, perplexity is 213.104\n",
      "After 800 steps, perplexity is 212.414\n",
      "After 900 steps, perplexity is 211.959\n",
      "After 1000 steps, perplexity is 210.966\n",
      "After 1100 steps, perplexity is 210.475\n",
      "After 1200 steps, perplexity is 209.590\n",
      "After 1300 steps, perplexity is 200.107\n",
      "After 1400 steps, perplexity is 200.050\n",
      "After 1500 steps, perplexity is 193.237\n",
      "After 1600 steps, perplexity is 192.123\n",
      "After 1700 steps, perplexity is 180.731\n",
      "After 1800 steps, perplexity is 183.199\n",
      "Epoch: 3 Validation Perplexity: 169.009\n",
      "Test Perplexity: 142.681\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "    train_data, valid_data, test_data, _ = ptb_raw_data(data_path)\n",
    "\n",
    "    # 计算一个epoch需要训练的次数\n",
    "    train_data_len = len(train_data)\n",
    "    train_batch_len = train_data_len // train_batch_size\n",
    "    train_epoch_size = (train_batch_len - 1) // train_num_step\n",
    "\n",
    "    valid_data_len = len(valid_data)\n",
    "    valid_batch_len = valid_data_len // eval_batch_size\n",
    "    valid_epoch_size = (valid_batch_len - 1) // eval_num_step\n",
    "\n",
    "    test_data_len = len(test_data)\n",
    "    test_batch_len = test_data_len // eval_batch_size\n",
    "    test_epoch_size = (test_batch_len - 1) // eval_num_step\n",
    "\n",
    "    initializer = tf.random_uniform_initializer(-0.05, 0.05)\n",
    "    with tf.variable_scope(\"language_model\", reuse=None, initializer=initializer):\n",
    "        train_model = PTBModel(True, train_batch_size, train_num_step)\n",
    "\n",
    "    with tf.variable_scope(\"language_model\", reuse=True, initializer=initializer):\n",
    "        eval_model = PTBModel(False, eval_batch_size, eval_num_step)\n",
    "\n",
    "    # 训练模型。\n",
    "    with tf.Session() as session:\n",
    "        tf.global_variables_initializer().run()\n",
    "\n",
    "        train_queue = ptb_producer(train_data, train_model.batch_size, train_model.num_steps)\n",
    "        eval_queue = ptb_producer(valid_data, eval_model.batch_size, eval_model.num_steps)\n",
    "        test_queue = ptb_producer(test_data, eval_model.batch_size, eval_model.num_steps)\n",
    "\n",
    "        coord = tf.train.Coordinator()\n",
    "        threads = tf.train.start_queue_runners(sess=session, coord=coord)\n",
    "\n",
    "        for i in range(num_epoch):\n",
    "            print(\"In iteration: %d\" % (i + 1))\n",
    "            run_epoch(session, train_model, train_queue, train_model.train_op, True, train_epoch_size)\n",
    "\n",
    "            valid_perplexity = run_epoch(session, eval_model, eval_queue, tf.no_op(), False, valid_epoch_size)\n",
    "            print(\"Epoch: %d Validation Perplexity: %.3f\" % (i + 1, valid_perplexity))\n",
    "\n",
    "        test_perplexity = run_epoch(session, eval_model, test_queue, tf.no_op(), False, test_epoch_size)\n",
    "        print(\"Test Perplexity: %.3f\" % test_perplexity)\n",
    "\n",
    "        coord.request_stop()\n",
    "        coord.join(threads)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
